package huffman;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.HashMap;

import bitstreams.BitInputStream;
import bitstreams.BitOutputStream;

public class HuffmanModel
{
	//alphabet
	public final char alphabet[];
	private HashMap<Character, Integer> alphabetIndexMap;	
	private HashMap<Character, TreeNode> transmittedMap;			
	private HashMap<TreeNode, Character> charByNode;			
	private ArrayList<TreeNode> nodesArr;	
	private TreeNode root;
	private TreeNode notYetTransmitted;
	
	 //alphabet length = 2^e + r
	public final int e;
	public final int r;
	
	public TreeNode getRoot() {
		return root;
	}
	
	 //LAST LETTER IN GIVEN ALPHABET WILL BE USED AS END OF INPUT!

	public HuffmanModel( char alphabet[] )
	{
		if( alphabet==null )
			throw new NullPointerException("Exception:accepted null alphabet!");
		
		e = (int) (Math.log(alphabet.length) / Math.log(2));
		r = alphabet.length - (int)Math.pow(2,e);
		this.alphabet = alphabet;
		
		alphabetIndexMap = new HashMap<Character, Integer>();
		
		for( int i=0 ; i<alphabet.length ; i++ )
		{
			alphabetIndexMap.put(alphabet[i], i);
		}
		
		reset();
	}
	
	public void reset()
	{
		this.notYetTransmitted = new TreeNode(alphabet.length*2 - 1);
		this.root = this.notYetTransmitted;
		this.transmittedMap = new HashMap<Character, TreeNode>();
		this.charByNode = new HashMap<TreeNode, Character>();
		this.nodesArr = new ArrayList<TreeNode>();
		this.nodesArr.add(this.root);
	}
	
	
	public void update(char c)
	{
		Integer index = alphabetIndexMap.get(c);
		
		if(index==null)
			throw new IllegalArgumentException("character '"+c+"' (code="+(int)c+") is not a part if the alphabet");
		
		TreeNode n = transmittedMap.get(c);
		
		if( n==null ) // first appearance of the character
		{
			n = splitNYT(c);
			
			if( n==root )
				return;
			else
				n = n.getParent();
		}
		else if( n.getParent()==this.notYetTransmitted.getParent() )
		{
			TreeNode highest = findMaximumLeafInBlock(n.getWeight());
			
			if( n!=highest )
				swapNodes(n, highest);
			
			n.incrementWeight();
			n = n.getParent();
		}
		

		while( n!=root )
		{
			TreeNode highest = findMaximumNodeInBlock(n.getWeight());
			
			if( n!=highest )
				swapNodes(n, highest);
			
			n.incrementWeight();
			n = n.getParent();
		}
		n.incrementWeight();
	}
	

	private TreeNode splitNYT(char c)
	{
		TreeNode newNYT = new TreeNode(this.notYetTransmitted.getNodeNumber()-2);
		TreeNode newExternal = new TreeNode(this.notYetTransmitted.getNodeNumber()-1);
		
		this.notYetTransmitted.setLeft(newNYT);
		this.notYetTransmitted.setRight(newExternal);
		
		TreeNode oldNYT = this.notYetTransmitted;
		this.notYetTransmitted = newNYT;
		
		oldNYT.incrementWeight();
		newExternal.incrementWeight();
		
		this.transmittedMap.put(c, newExternal);
		this.charByNode.put(newExternal, c);
		this.nodesArr.add(newExternal);
		this.nodesArr.add(newNYT);

		return oldNYT;
	}
	
	/**
	 * search for the node with the highest node number who's weight equals 'weight'
	 */
	private TreeNode findMaximumNodeInBlock( int weight )
	{
		TreeNode maxnode = null;
		for( TreeNode n : nodesArr )
		{
			if(n!=root)
			{
				if( maxnode==null && n.getWeight()==weight ) 
					maxnode = n;
				else if( n.getWeight()==weight )
				{
					if( maxnode.getNodeNumber() < n.getNodeNumber() )
						maxnode = n;
				}
			}
		}
		
		return maxnode;
	}
	
	/**
	 * search for the leaf with the highest node number who's weight equals 'weight'
	 */
	private TreeNode findMaximumLeafInBlock( int weight )
	{
		TreeNode maxnode = null;
		for( TreeNode n : nodesArr )
		{
			if(n!=root && n.isLeaf())
			{
				if( maxnode==null && n.getWeight()==weight ) 
					maxnode = n;
				else if( n.getWeight()==weight )
				{
					if( maxnode.getNodeNumber() < n.getNodeNumber() )
						maxnode = n;
				}
			}
		}
		
		return maxnode;
	}
	
	private void swapNodes( TreeNode a, TreeNode b )
	{
		TreeNode parentOf_a = a.getParent();
		TreeNode parentOf_b = b.getParent();
		
		boolean aWasLeft = a.isLeftChild();
		boolean bWasLeft = b.isLeftChild();
		
		a.detachFromParent();
		b.detachFromParent();
		
		if( aWasLeft )
			parentOf_a.setLeft(b);
		else
			parentOf_a.setRight(b);
		
		if( bWasLeft )
			parentOf_b.setLeft(a);
		else
			parentOf_b.setRight(a);
		
		int temp = a.getNodeNumber();
		a.setNodeNumber(b.getNodeNumber());
		b.setNodeNumber(temp);
	}
	
	public CodeParams getCode(char c)
	{
		TreeNode n = transmittedMap.get(c);
		
		CodeParams p = null;
		
		if(n==null)
		{
			p = getDefaultCode(c);
			n = this.notYetTransmitted;
		}
		else
			p = new CodeParams();
		
		while( n!=root )
		{
			if(n.isRightChild())
				p.code |= 1<<p.length;
			
			p.length++;
			
			n = n.getParent();
		}
		
		return p;
	}
	
	public CodeParams getDefaultCode(char c)
	{
		Integer index = alphabetIndexMap.get(c);
		
		if(index==null)
			throw new IllegalArgumentException("character '"+c+"' ("+(int)c+") is not a part if the alphabet");
		
		CodeParams p = new CodeParams();
		
		if( index < this.r*2 )
		{
			p.code = index;
			p.length = this.e + 1;
		}
		else
		{
			p.code = index - this.r;
			p.length = this.e;			
		}
		return p;
	}
	
	/**
	 * decode input using adaptive Huffman coding algorithm.
	 * @throws IOException
	 */
	public void decode( InputStream inStream, OutputStream outStream ) throws IOException
	{
		BitInputStream input = new BitInputStream(inStream);
		BufferedWriter outWriter = new BufferedWriter(new OutputStreamWriter(outStream));
		
		TreeNode curTravesal = this.root;
		int bit = 0;
		
		while( bit!=-1 )
		{
			if( curTravesal.isLeaf() )
			{
				char c ;
				if( curTravesal==this.notYetTransmitted )
				{
					int b = input.read(this.e);
					
					if( b<this.r )
						b = b<<1 | input.read(1);
					
					else
						b += this.r;
					
					c = this.alphabet[b];
				}
				else
					c = this.charByNode.get(curTravesal);
				
				if(c == alphabet[alphabet.length-1] )
					break;
				
				outWriter.write(c);
				this.update(c);
				curTravesal = this.root;
			}
			
			bit = input.read(1);
			
			if(bit==0)
				curTravesal = curTravesal.getLeft();
			else
				curTravesal = curTravesal.getRight();
		}
		
		input.close();
		outWriter.close();
	}
	
	public void encode(InputStream inStream, OutputStream outStream) throws IOException
	{
		BitOutputStream output = new BitOutputStream(outStream);
		
		BufferedReader input = new BufferedReader( new InputStreamReader(inStream) );
		
		char c;
		while( (c=(char) input.read()) != 65535 )
		{
			CodeParams p = this.getCode(c);
			output.write(p.length, p.code);
			this.update(c);
		}
		
		CodeParams p = this.getCode(alphabet[alphabet.length-1]);
		output.write(p.length, p.code);
		this.update(alphabet[alphabet.length-1]);
		
		input.close();
		output.close();
	}
}
